import win32pdh

class PerfInfo:
    def __init__(self, counter_path):
        self.path = counter_path
        if win32pdh.ValidatePath(self.path) != 0:
            raise Exception('Invalid path: %s' % counter_path)
        self.query = self.counter = None

    def __del__(self):
        self.stop()

    def start(self):
        try:
            self.query = win32pdh.OpenQuery(None, 0)
            self.counter = win32pdh.AddCounter(self.query, self.path, 0)
        except:
            self.stop()
            raise

    def stop(self):
        if self.counter is not None:
            win32pdh.RemoveCounter(self.counter)
        if self.query is not None:
            win32pdh.CloseQuery(self.query)
        self.counter = self.query = None

    def get_value(self):
        win32pdh.CollectQueryData(self.query)
        #print self.counter, win32pdh.PDH_FMT_LONG
        status, value = win32pdh.GetFormattedCounterValue(self.counter, win32pdh.PDH_FMT_LONG)
        return value
            

class CPUUsage(PerfInfo):
    def __init__(self):
        PerfInfo.__init__(self, r'\Processor(_Total)\% Processor Time')

    def start(self):
        # Ignore the first collected data because it always is 99%
        PerfInfo.start(self)
        self.get_value()
        
class DiskUsage(PerfInfo):
    def __init__(self):
        PerfInfo.__init__(self, r'\PhysicalDisk(_Total)\% Disk Time')

    def start(self):
        # Ignore the first collected data because it always is 99%
        PerfInfo.start(self)
        self.get_value()


def _test():
    import time

    cpu_usage = CPUUsage()
    cpu_usage.start()

    try:
        while True:
            percent = cpu_usage.get_value()
            print 'CPU Usage: %s (Press Ctrl-C to stop)' % percent
            time.sleep(1)
    except KeyboardInterrupt:
        pass
        
    cpu_usage.stop()
    
def _test2():
    import time

    disk_usage = DiskUsage()
    disk_usage.start()

    try:
        while True:
            percent = disk_usage.get_value()
            print 'Disk Usage: %s (Press Ctrl-C to stop)' % percent
            time.sleep(1)
    except KeyboardInterrupt:
        pass
        
    disk_usage.stop()

if __name__ == '__main__':
    _test2()